from process_data import prepare_splits
from model import CNN
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import torch

DATA_PATH = "../MERGED_ECSMP_BVP"
RANDOM_SEED = 42
EPOCHS = 1
LEARNING_RATE = 0.001
#LEARNING_RATE = 0.0001

def train_model(model, train_loader, test_loader, device, epochs=EPOCHS):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=2, verbose=True)

    best_accuracy = 0.0

    for epoch in range(epochs):
        # Training phase
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}")
        for inputs, labels in progress_bar:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            # Update progress bar
            progress_bar.set_postfix({
                'loss': running_loss / (progress_bar.n + 1),
                'acc': 100 * correct / total
            })

        train_accuracy = 100 * correct / total

        # Evaluation phase
        model.eval()
        test_loss = 0
        test_correct = 0
        test_total = 0

        with torch.no_grad():
            for inputs, labels in tqdm(test_loader, desc="Testing"):
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                test_loss += loss.item()

                _, predicted = torch.max(outputs.data, 1)
                test_total += labels.size(0)
                test_correct += (predicted == labels).sum().item()

        test_accuracy = 100 * test_correct / test_total

        # Adjust learning rate based on validation accuracy
        scheduler.step(test_accuracy)

        print(f"Epoch {epoch + 1}/{epochs}")
        print(f"Train Loss: {running_loss / len(train_loader):.4f}, Train Accuracy: {train_accuracy:.2f}%")
        print(f"Test Loss: {test_loss / len(test_loader):.4f}, Test Accuracy: {test_accuracy:.2f}%")

        # Save the best model
        if test_accuracy > best_accuracy:
            best_accuracy = test_accuracy
            torch.save(model.state_dict(), "./best_user_identification_model.pth")
            print(f"New best model saved with accuracy: {best_accuracy:.2f}%")

    return model, best_accuracy


def evaluate_model(model, test_loader, device):
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for inputs, labels in tqdm(test_loader, desc="Final Evaluation"):
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Final Test Accuracy: {accuracy:.2f}%")

    return accuracy

if __name__ == '__main__':

    #torch.manual_seed(RANDOM_SEED)
    #np.random.seed(RANDOM_SEED)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device {device}")

    # Prepare data
    train_loader, test_loader, num_classes = prepare_splits()

    # Get sample input shape from first batch
    for inputs, _ in train_loader:
        input_shape = inputs.shape
        break

    print(f"Input shape: {input_shape}")

    # Initialize the model
    model = CNN(num_classes=num_classes).to(device)
    print(f"Model initialized with {num_classes} classes")

    # Train the model
    model, best_accuracy = train_model(model, train_loader, test_loader, device, EPOCHS)

    # Load the best model for final evaluation
    model.load_state_dict(torch.load("best_user_identification_model.pth"))

    # Final evaluation
    final_accuracy = evaluate_model(model, test_loader, device)

    print(f"Training completed. Best accuracy: {best_accuracy:.2f}%")
    print(f"Model saved as 'best_user_identification_model.pth'")